;; Multiarray implementation for Unidata array class

;; This module requires the ParallelColt library (version 0.9)
;; on the classpath.

;; by Konrad Hinsen
;; last updated February 8, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(use 'nstools.ns)
(ns+ multiarray.colt
  (:clone multiarray.utility)
  (:use [multiarray.protocol]))

;;
;; Array constructors
;;

(defvar- factory-for
  {Integer/TYPE [cern.colt.matrix.tint.IntFactory1D/dense
		 cern.colt.matrix.tint.IntFactory2D/dense
		 cern.colt.matrix.tint.IntFactory3D/dense]
   Long/TYPE    [cern.colt.matrix.tlong.LongFactory1D/dense
		 cern.colt.matrix.tlong.LongFactory2D/dense
		 cern.colt.matrix.tlong.LongFactory3D/dense]
   Float/TYPE   [cern.colt.matrix.tfloat.FloatFactory1D/dense
		 cern.colt.matrix.tfloat.FloatFactory2D/dense
		 cern.colt.matrix.tfloat.FloatFactory3D/dense]
   Double/TYPE  [cern.colt.matrix.tdouble.DoubleFactory1D/dense
		 cern.colt.matrix.tdouble.DoubleFactory2D/dense
		 cern.colt.matrix.tdouble.DoubleFactory3D/dense]
   Object       [cern.colt.matrix.tobject.ObjectFactory1D/dense
		 cern.colt.matrix.tobject.ObjectFactory2D/dense
		 cern.colt.matrix.tobject.ObjectFactory3D/dense]})

(defn- array-element-class
  [datatype]
  (let [type-map {Boolean Integer/TYPE Boolean/TYPE Integer/TYPE
                  Byte    Integer/TYPE Byte/TYPE Integer/TYPE
		  Short   Integer/TYPE Short/TYPE Integer/TYPE
		  Integer Integer/TYPE Integer/TYPE Integer/TYPE
                  Long    Long/TYPE    Long/TYPE Long/TYPE
                  Float   Float/TYPE   Float/TYPE Float/TYPE
                  Double  Double/TYPE  Double/TYPE Double/TYPE}]
    (get type-map datatype Object)))

; zeros

(defn zeros
  ([shape]
     (zeros shape Integer/TYPE))
  ([shape datatype]
     (let [n (count shape)
	   f (get (factory-for (array-element-class datatype)) (dec n))]
       (cond (= n 1) (.make f (get shape 0))
     	     (= n 2) (.make f (get shape 0) (get shape 1))
     	     (= n 3) (.make f (get shape 0) (get shape 1) (get shape 2))
	     :else   (throw (Exception.
			     "Colt arrays can only have rank 1, 2, or 3."))))))

;;
;; Implementation of clj-multiarray.protocol
;;

(defvar- lowlevel-methods
  {:element-class
   (fn [#^cern.colt.matrix.AbstractMatrix ma]
     ; the call to elements uses reflection
     (.. ma (elements) (getClass) (getComponentType)))
   :element-array
   (fn [#^cern.colt.matrix.AbstractMatrix ma]
     ; the call to elements uses reflection
     (.elements ma))})

(extend cern.colt.matrix.AbstractMatrix1D
  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape (fn [#^cern.colt.matrix.AbstractMatrix1D ma]
		   [(.size ma)])
	  :rank  (fn [#^cern.colt.matrix.AbstractMatrix1D ma]
		   1)})
  MultiArrayLowLevelProtocol
    lowlevel-methods)

(extend cern.colt.matrix.AbstractMatrix2D
  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape (fn [#^cern.colt.matrix.AbstractMatrix2D ma]
		   [(.rows ma) (.columns ma)])
	  :rank  (fn [#^cern.colt.matrix.AbstractMatrix2D ma]
		   2)})
  MultiArrayLowLevelProtocol
    lowlevel-methods)
            
(extend cern.colt.matrix.AbstractMatrix3D
  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape (fn [#^cern.colt.matrix.AbstractMatrix3D ma]
		   [(.slices ma) (.rows ma) (.columns ma)])
	  :rank  (fn [#^cern.colt.matrix.AbstractMatrix3D ma]
		   3)})
  MultiArrayLowLevelProtocol
    lowlevel-methods)

;;
;; Conversion from other multiarray implementations
;;
(defn- into-ma-1
  [[n] array]
  (let [element-type (.getComponentType (class array))]
    (cond (isa? element-type Object)
	    (cern.colt.matrix.tobject.impl.DenseObjectMatrix1D. array)
          (identical? element-type Integer/TYPE)
	    (cern.colt.matrix.tint.impl.DenseIntMatrix1D.
	       n array 0 1 false)
	  (identical? element-type Long/TYPE)
	    (cern.colt.matrix.tlong.impl.DenseLongMatrix1D.
	       n array 0 1 false)
	  (identical? element-type Float/TYPE)
	    (cern.colt.matrix.tfloat.impl.DenseFloatMatrix1D.
	       n array 0 1 false)
	  (identical? element-type Double/TYPE)
	    (cern.colt.matrix.tdouble.impl.DenseDoubleMatrix1D.
	       n array 0 1 false))))

(defn- into-ma-2
  [[rows cols] array]
  (let [element-type (.getComponentType (class array))]
    (cond (isa? element-type Object)
	    (throw (Exception. "not implemented"))
          (identical? element-type Integer/TYPE)
	    (cern.colt.matrix.tint.impl.DenseIntMatrix2D.
	       rows cols array 0 0 cols 1 false)
	  (identical? element-type Long/TYPE)
	    (cern.colt.matrix.tlong.impl.DenseLongMatrix2D.
	       rows cols array 0 0 cols 1 false)
	  (identical? element-type Float/TYPE)
	    (cern.colt.matrix.tfloat.impl.DenseFloatMatrix2D.
	       rows cols array 0 0 cols 1 false)
	  (identical? element-type Double/TYPE)
	    (cern.colt.matrix.tdouble.impl.DenseDoubleMatrix2D.
	       rows cols array 0 0 cols 1 false))))

(defn- into-ma-3
  [[slices rows cols] array]
  (let [element-type (.getComponentType (class array))]
    (cond (isa? element-type Object)
	    (throw (Exception. "not implemented"))
          (identical? element-type Integer/TYPE)
	    (cern.colt.matrix.tint.impl.DenseIntMatrix3D.
	       slices rows cols array 0 0 0 (* rows cols) cols 1 false)
	  (identical? element-type Long/TYPE)
	    (cern.colt.matrix.tlong.impl.DenseLongMatrix3D.
	       slices rows cols array 0 0 0 (* rows cols) cols 1 false)
	  (identical? element-type Float/TYPE)
	    (cern.colt.matrix.tfloat.impl.DenseFloatMatrix3D.
	       slices rows cols array 0 0 0 (* rows cols) cols 1 false)
	  (identical? element-type Double/TYPE)
	    (cern.colt.matrix.tdouble.impl.DenseDoubleMatrix3D.
	       slices rows cols array 0 0 0 (* rows cols) cols 1 false))))

(defn into-ma
  [obj]
  (let [sh    (shape obj)
	n     (count sh)
	array (element-array obj)]
    (cond (= n 1) (into-ma-1 sh array)
	  (= n 2) (into-ma-2 sh array)
	  (= n 3) (into-ma-3 sh array)
	  :else (throw (Exception.
			"Colt matrices exist for ranks 1, 2, 3 only")))))
